"""
进度推送API端点
"""
import logging
from typing import Dict, Any
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from sqlalchemy.orm import Session
from app.db.session import get_db
from app.api.v1.endpoints.websocket import websocket_manager
from app.models.video_task import VideoTask
from uuid import UUID

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/ws/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str):
    """
    WebSocket端点，用于实时推送任务进度
    """
    try:
        # 验证任务ID格式
        UUID(task_id)
    except ValueError:
        await websocket.close(code=4000, reason="Invalid task ID")
        return

    connection_id = None

    try:
        # 建立连接
        connection_id = await websocket_manager.connect(websocket, task_id)

        # 保持连接活跃
        while True:
            try:
                # 等待客户端消息（心跳包）
                data = await websocket.receive_text()
                message = {"type": "ping", "data": data}
                await websocket_manager.send_to_connection(connection_id, {
                    "type": "pong",
                    "task_id": task_id,
                    "message": "心跳响应"
                })

            except WebSocketDisconnect:
                break
            except Exception as e:
                logger.error(f"WebSocket error for {connection_id}: {e}")
                break

    except WebSocketDisconnect:
        logger.info(f"WebSocket disconnected for task {task_id}")
    except Exception as e:
        logger.error(f"WebSocket connection error: {e}")
    finally:
        if connection_id:
            websocket_manager.disconnect(connection_id)


@router.post("/notify/{task_id}")
async def notify_progress(
    task_id: str,
    progress_data: Dict[str, Any],
    db: Session = Depends(get_db)
):
    """
    手动触发进度通知（用于Celery任务）
    """
    try:
        # 验证任务是否存在
        task = db.query(VideoTask).filter(VideoTask.id == task_id).first()
        if not task:
            return {"error": "Task not found"}

        # 更新任务状态
        if "status" in progress_data:
            task.status = progress_data["status"]
        if "progress" in progress_data:
            task.progress = progress_data["progress"]
        if "current_step" in progress_data:
            task.current_step = progress_data["current_step"]
        if "error_message" in progress_data:
            task.error_message = progress_data["error_message"]

        db.commit()

        # 发送WebSocket通知
        message = {
            "type": "progress_update",
            "task_id": task_id,
            "data": progress_data,
            "timestamp": str(task.updated_at) if task.updated_at else None
        }

        await websocket_manager.send_to_task(task_id, message)

        return {"success": True, "message": "Progress notification sent"}

    except Exception as e:
        logger.error(f"Failed to send progress notification: {e}")
        return {"error": str(e)}


@router.get("/status/{task_id}")
async def get_progress_status(task_id: str, db: Session = Depends(get_db)):
    """
    获取任务当前进度状态
    """
    try:
        task = db.query(VideoTask).filter(VideoTask.id == task_id).first()
        if not task:
            return {"error": "Task not found"}

        # 检查是否有活跃的WebSocket连接
        active_connections = websocket_manager.get_connection_count(task_id)

        return {
            "task_id": task_id,
            "status": task.status,
            "progress": task.progress,
            "current_step": task.current_step,
            "error_message": task.error_message,
            "active_connections": active_connections,
            "updated_at": task.updated_at,
            "created_at": task.created_at
        }

    except Exception as e:
        logger.error(f"Failed to get progress status: {e}")
        return {"error": str(e)}


@router.get("/stats")
async def get_websocket_stats():
    """
    获取WebSocket连接统计信息
    """
    try:
        active_tasks = websocket_manager.get_active_tasks()
        total_connections = websocket_manager.get_connection_count()

        task_details = {}
        for task_id in active_tasks:
            task_details[task_id] = websocket_manager.get_connection_count(task_id)

        return {
            "total_connections": total_connections,
            "active_tasks": len(active_tasks),
            "task_details": task_details
        }

    except Exception as e:
        logger.error(f"Failed to get WebSocket stats: {e}")
        return {"error": str(e)}